[None][feat] Load static FP8 Cosmos3 Nano/Super checkpoints without re-quantizing them - #17476

Dev Engineer Review

• Added static ModelOpt FP8 loading for single-GPU Cosmos3 Nano and Super checkpoints.
• Preserved separate Q/K/V and gate/up projections for static FP8 checkpoint fidelity.
• Added shared QKV input quantization and split GatedMLP support.
• Added silu_and_mul_2in and swiglu_2in with FP8 output support and B200 tuning.
• Preserved fused topology for BF16, dynamic FP8, and other models.
• Added BF16 activation execution for the denoising steps the checkpoint's diffusion_step_policy nominates.
• Added destination-dtype validation for static scale loading.
• Rejected unsupported multi-GPU static FP8 configurations and unsupported split GatedMLP options.
• Added documentation and checkpoint-fidelity coverage for Cosmos3 Nano/Super FP8.
• CI failed without individual test results. Re-trigger CI before final approval.

QA Engineer Review

Test changes

Added tests for:

• test_swiglu_2in.py: two-input SwiGLU behavior, quantization, compilation, CUDA graphs, validation, and B200 tuning.
• test_gated_mlp_split.py: split and fused topology, checkpoint loading, quantization, compilation, CUDA graphs, scale validation, BF16, and unsupported configurations.
• test_attention_qkv_share.py: shared quantization, validation, numerical equivalence, bypass paths, and rank handling.
• test_cosmos3_fp8.py: static and dynamic FP8 resolution, module layouts, topology selection, scale validation, and exact checkpoint loading.
• test_quant_static_guard.py: high-precision destination handling and incompatible FP8 destination rejection.
• test_cosmos3_step_precision.py and test_cosmos3_step_precision_component.py: step-window selection, precision dispatch, activation quantization behavior, reset and validation, and FP8 component accuracy.
• test_visual_gen_cosmos3.py: static FP8 T2I generation, topology validation, and LPIPS regression.

Test-list coverage

tests/integration/test_lists/test-db/l0_b200.yml includes entries for the SwiGLU, split GatedMLP, shared-QKV, Cosmos3 FP8, step-precision component, and Cosmos3 Nano FP8 LPIPS tests.

tests/integration/test_lists/test-db/l0_cpu.yml includes test_cosmos3_step_precision.py.

The static quantization guard test has no reported test-list entry.

Verdict: needs follow-up because CI provides no individual test results and complete coverage cannot be confirmed until CI is re-triggered.

Description

The pre-quantized (ModelOpt-calibrated) Cosmos3 Nano and Super FP8 checkpoints
already loaded, but they did not load faithfully.

Static FP8 calibrates one scale per projection. The Cosmos3 topology fused the
generation tower's Q/K/V into one Linear and both towers' gate/up into
another. A per-tensor FP8 Linear holds exactly one weight scale, so the loader
kept max(scales) and re-quantized the other members of each group onto it:

qWi' = FP8-round((si * qWi) / max(s)) instead of keeping si * qWi

That is a second rounding on top of ModelOpt's. It affected 141 of 252 fused
shards in Nano and 252 of 448 in Super (~56% in both), and on affected shards
the added error was 7-18x the numerical floor measured without the extra
conversion. The checkpoints were fine; the fusion was discarding their
calibration.

This PR keeps those projections separate under static FP8, so every tensor and
scale loads exactly as calibrated. One predicate,
transformer_cosmos3.uses_static_fp8(), gates all of it. BF16 Cosmos3,
dynamically quantized Cosmos3, and every other model keep the fused topology
unchanged.

vLLM-Omni's Cosmos3 implementation also keeps these projections separate; the
single-scale fusion was specific to this topology.

It is also faster

Unfusing was expected to cost a little. It does not — the split topology is
2.13% faster end to end (Nano, default T2V, one B200: 104.62 s vs 106.90 s,
variants interleaved at the load level, 4 measured runs each after a discarded
warmup, ranges non-overlapping).

The win is not in the GEMMs (CUPTI puts split QKV at 1.03x fused and gate/up at
1.00x). Fusing QKV hands downstream code qkv.split() views that keep the
fused row stride, so per-head QK RMSNorm, RoPE and the und/gen K/V concatenation
all read strided. Split projections are contiguous Linear outputs. A
kernel-level profile diff attributes the delta to exactly those ops.

What changed

Shared infrastructure (opt-in, inert for every existing caller):

• swiglu.py / torch_custom_ops.py — a new two-input SwiGLU,
  trtllm::silu_and_mul_2in, computing SiLU(G) * U from two separate tensors
  and quantizing the result directly for the FP8 down_proj. Bit-exact against
  the shipped fused kernel on every tested BF16/FP16 shape, with swiglu_limit,
  with swiglu_alpha/swiglu_beta, and when emitting FP8. It is 0.79x the
  shipped fused activation on the FP8 path Cosmos3 takes (1,047 -> 824 us Nano,
  2,211 -> 1,725 us Super), DRAM utilization 68% -> 86-87% (B200). Full
  custom_op contract: register_fake, opcheck,
  torch.compile(fullgraph=True), CUDA-graph capture.

  Its launch parameters (get_silu_b200_tuning_params) are constants from a
  B200 sweep, deliberately scoped rather than generalized: a B300 (sm_103)
  sweep puts them within 0.4% of the best configuration measured there, and the
  docstring directs a future architecture to add its own table and dispatch on
  device capability rather than widening this one.

  torch.cat([gate, up]) was rejected — it adds a full intermediate-sized GPU
  copy — as was writing both GEMMs into column slices of one [M, 2I] buffer,
  which is silently mis-written because cublas_gemm_caller hardcodes
  ldc = n (cublasScaledMM.cpp) while cublas_scaled_mm_out still validates
  the output stride.

• gated_mlp.py — opt-in GatedMLP(split_gate_up=...). Quantizes the shared
  activation once and hands both projections the FP8 tensor
  (FP8QDQLinearMethod.apply passes a pre-quantized input straight through).
  Rejects LoRA and force_dynamic_quantization explicitly rather than by
  accident.

• visual_gen/modules/attention.py — opt-in
  Attention(share_qkv_input_quant=...), the FP8 analog of the existing NVFP4
  dedup, with the same single-quantization property for Q/K/V.

Both carry a post_load_weights() validating the invariant the sharing rests
on — that the group's projections carry the same calibrated input_scale.
Checked at load, never in forward(): reading a scale tensor on the hot path
would sync the device and break fullgraph compilation. The invariant was
verified exhaustively across the shipped checkpoints — 108 groups in Nano, 192
in Super, zero with differing input scales — but a checkpoint that violated
it would otherwise produce silently wrong output.

Cosmos3 wiring (transformer_cosmos3.py): uses_static_fp8() drives all
four sites — GEN cross-attention (SEPARATE_QKV + shared quant), UND causal
attention (shared quant; already separate), and both towers' gate/up.
post_load_weights() gained a second pass over GatedMLP/Attention, because
named_modules() yields parents before children and the invariants must be
checked after the projections finalize.

Since Cosmos3-Edge (#16773) landed, both decoder layers build their MLP through
_build_cosmos3_mlp(), so split_gate_up=uses_static_fp8(model_config) is
passed there rather than at the two call sites. That is load-bearing: dropping
it leaves static FP8 on the fused topology, which still loads and still renders
plausible output while re-quantizing 56% of shards — the exact defect this
change removes. test_topology_follows_quantization and
test_static_fp8_checkpoint_realizes_expected_module_layout, which pins the
exact per-tower Linear counts, both fail if it regresses.

Edge's Nemotron-dense recipe (recipe.gated_mlp false) builds MLP instead,
which has no gate/up pair to split and is unaffected. A statically quantized
Edge checkpoint would need its own handling; none exists today.

The loader needed no change. DynamicLinearWeightLoader dispatches on each
module's own weight_mode, and the existing key remap already emits
per-projection checkpoint names, so a split (VANILLA) module matches its
checkpoint tensor directly while params_map continues to serve the fused BF16
path untouched.

Supported surface, and what is refused

Nano and Super; T2V, T2I, I2V and V2V; one GPU; static per-tensor
QuantAlgo.FP8 checkpoints only.

Static FP8 raises NotImplementedError for tp_size, ulysses_size,
cfg_size, cp_size or parallel_vae_size above 1, rather than running
untested. The guard is deliberate: putting cross-attention on SEPARATE_QKV
changes how the parallel wrappers treat it — with Attention2D engaged, Ulysses
is disabled rather than failing, announced only by a logger.debug — and a
quiet fallback is worse than a refusal. Existing distributed Cosmos3 tests use
synthetic BF16 weights and do not validate FP8 sharding, per-rank scales, or the
NVLink/NCCL collective paths.

Also out of scope, and documented as such in the example README:

• Audio (T2AV/TI2AV) runs, and is not refused. Both checkpoints ship the
  audio tower (sound_gen: true), so audio is quantized and generated like any
  other task; the single-GPU guard covers only the parallel axes. It is simply
  less exercised than the four video/image tasks and carries no quality claim.
• No quality measurement of our own. There is no accuracy or
  generative-quality benchmark for these checkpoints here. This PR claims
  functional support and checkpoint fidelity — not that FP8 and BF16 are
  visually equivalent. Step precision is enabled by default on the strength of
  the recipe authors' reports and two independent implementations, not on a
  measurement taken in this repository; the component tests establish that it
  reduces activation-quantization error at the component level, which is a
  narrower claim than image quality.
• No Hub IDs. The checkpoints live in subdirectories and are addressed by
  local path. hf_ids registration is deferred until directly loadable official
  FP8 repository IDs exist.

One interaction is documented but left unguarded: SEPARATE_QKV
cross-attention also falls back from CUTEDSL VSA to VANILLA, which would drop
sparse attention under FP8 with no diagnostic. No shipped Cosmos3 config enables
VSA, so the combination is unreachable today.

Test Coverage

New, 78 unit tests:

• tests/unittest/_torch/modules/test_swiglu_2in.py (32) — torch.library.opcheck,
  torch.compile(fullgraph=True), CUDA-graph capture, rank-N inputs,
  non-contiguous rejection, and bit-exactness against the shipped kernel for
  BF16 / FP16 / swiglu_limit / swiglu_alpha+swiglu_beta / FP8 output.
• tests/unittest/_torch/modules/test_gated_mlp_split.py (19) — split topology,
  single-quantization by storage identity (data_ptr, not dtype — three
  independent quantizations would also yield FP8 everywhere), scale-mismatch
  rejection, LoRA and dynamic-quantization refusals.
• tests/unittest/_torch/visual_gen/test_attention_qkv_share.py (12) — the same
  for Q/K/V, including the cross-attention case where k/v come from a different
  tensor and sharing must not engage.
• tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py (15) — checkpoint-gated.
  test_previously_fused_groups_now_load_exactly compares FP8 weights bitwise
  against the raw checkpoint tensors, and first asserts the group's weight
  scales genuinely differ — otherwise fusion would have been lossless and
  exactness would hold trivially under either topology, leaving the test unable
  to discriminate. test_topology_follows_quantization and
  test_dynamic_quantization_stays_fused pin all four directions so the BF16
  and dynamic paths cannot drift silently.

Regression: the existing test_fused_activation_quant.py and
test_cosmos3_transformer.py suites stay green.

Step precision (new, see below):

• tests/unittest/_torch/visual_gen/test_cosmos3_step_precision.py (24, CPU) —
  the step policy (which steps are selected, that selection is a pure function
  of the step index so CFG branches cannot disagree, the one-step warmup
  carve-out), dispatch between the two paths, the dequantization arithmetic,
  and the refusal to accept an already-quantized activation.
• tests/unittest/_torch/visual_gen/test_cosmos3_step_precision_component.py
  (7, GPU) — a real quantized GatedMLP driven through
  the checkpoint's declared policy, comparing feature-on against feature-off in the same
  job: that the shared-activation optimization stands down on edge steps and
  stays engaged on middle ones, that the two paths produce different output at
  all, and that the 16-bit step lands closer to the unquantized reference
  than the quantized step. No stored artifact — the comparisons are exact
  arithmetic or an A/B against the module's other path.

No LPIPS golden. An earlier revision of this PR added one; it has been
removed. It rendered T2I at 4 steps, which the step-precision default puts
entirely on the 16-bit path, so the committed image no longer matched the path
under test. Re-cutting was not worth it: it was a per-checkpoint reference cut
on B300 (sm_103) but scheduled only on B200 (sm_100), and a stored quantized
golden does not survive that boundary — kernel selection, split-k and reduction
order and autotuner choices all change with the architecture. The checkpoints
are not staged in CI, so it skipped there in any case. What it protected is
covered by the bitwise weight/scale comparison and the topology-count assertions
in test_cosmos3_fp8.py, which are deterministic and architecture-independent.

CI (l0_b200.yml, l0_cpu.yml): the checkpoint-free suites run pre-merge
beside the other visual_gen unit suites. The FP8 checkpoints are not staged in
the CI llm-models/ tree, so test_cosmos3_fp8.py skips there — a skip
reports green, so it should not be read as evidence that the FP8 path is
exercised in CI today. It is exercised locally: 15/15 against both the Nano and
Super checkpoints, including the published Cosmos3-Super-Text2Image FP8 build.

Denoising-step activation precision

A ModelOpt checkpoint carries one activation scale per projection, calibrated as
a max over the whole sampling trajectory. That single scale fits the first and
last denoising steps worst, and the recipe's authors attribute frame-to-frame
flicker in video output to it. Those steps now run the resident FP8 weights
through a 16-bit GEMM instead — the weight is dequantized with its own
weight_scale and input_scale goes unused — while the middle steps keep the
checkpoint's fully quantized path.

No extra weights are read: same weights, same scales, no second checkpoint and
no persistent dequantized copy.

The recipe comes from the checkpoint. The producer publishes it under
quantization_config.runtime.diffusion_step_policy — which steps take the
16-bit path, and what the understanding tower does — and ships it only with the
builds whose calibration needs it. Of the six published FP8 builds, nano,
super and super-i2v carry a policy; the image build and both distilled
4-step builds deliberately do not, because the flicker this targets is not
observed there. A checkpoint without a policy therefore runs fully quantized,
and this PR exposes no configuration of its own: there is nothing left for a
knob to decide, and enabling the feature for a checkpoint that did not ask for
it would be wrong for half the fleet. It also means distilled schedules need no
special-casing.

Policy shapes we do not implement are refused rather than partially honoured:
unknown or missing fields, schema_version other than 1, and any other type,
index_space, default_mode or overlap value. Silently ignoring a field the
producer set is indistinguishable from the feature not working. overlap is
validated and not acted on — it selects which mode wins where the two windows
meet, which the predicate already yields; vLLM-Omni treats it the same way.

vLLM-Omni (vllm-project/vllm-omni#6560) and SGLang (sgl-project/sglang#36380)
landed the same mitigation independently, with the same one-step warmup
carve-out. This matches their semantics.

The reasoner is not step-gated. The policy states its precision outright,
because the understanding tower runs once per request — on whichever
transformer call builds its KV cache — so deriving it from a step index matches
the published 3-and-3 policy only because step 0 falls inside the first window,
and would stop matching for a policy with first_steps: 0. vLLM-Omni resolves
it the same way.

Unlike either implementation, this topology quantizes shared activations
above the Linear: gate/up and q/k/v each quantize once and hand the same
tensor to their projections, and swiglu_2in emits FP8 straight into
down_proj. All three must stand down while a 16-bit step is selected, or the
step still runs on FP8 activations and the feature is silently absent. A
quantization method advertises that by publishing high_precision; the sharing
sites consult it, and apply_fp8_w8a16_linear raises rather than accept an
already-quantized activation.

Coverage note: the component tests install the wrapper themselves, so they
cannot see the transformer reading the wrong config key or handing the
unconditional path to the wrong tower. Transformer-level wiring tests cover
both, and both mutations were confirmed to fail them.

A loader fix this path depends on

#17699 added a guard refusing static-quant recipes against unquantized
checkpoints. It resolves the algorithm by module name, falling back to the
global recipe for any module without its own quant_config — which includes
modules that cannot be quantized at all. Embedding reaches the quantized
linear loader because it subclasses LMHead -> Linear, yet its __init__
never exposes quant_config, so it always keeps a high-precision buffer, and
ModelOpt does not list it in ignore because only Linear targets were ever
candidates. Every static-FP8 VisualGen checkpoint therefore failed to load on
language_model.embed_tokens.

The guard now consults the destination buffer, which is the condition its own
docstring describes. A module built for FP8 holds a float8 buffer and still
raises; where the destination is unknown the check proceeds, keeping the
fail-closed behaviour. Both directions are pinned in
test_quant_static_guard.py. Static FP8 is the only pre-quantized recipe in the
tree, so this was latent until this PR.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.